import multiprocessing
from pprint import pprint

import pandas as pd
from cadCAD import Experiment, configs
from cadCAD.configuration.utils import config_sim
from cadCAD.engine import ExecutionMode, ExecutionContext, Executor
from matplotlib import pyplot as plt
from pyDOE import pbdesign, bbdesign
import numpy as np

init_state = {
    'height': 0,
    'leaf_num': 0
}
params = {
    'N': 1,
    'T': 10,
    'sunlight': ['low', 'high'],
    'water': ['low', 'high'],
    'temperature': ['low', 'high']
}


def getDesign(param, flag):
    design = None
    # design1: 生成4个因子的Plackett-Burman设计表(默认2水平)
    if flag == 1:
        print("设计方法： Plackett-Burman设计")
        design = (pbdesign(param) + 1) / 2
    # design2: 生成4个因子的Box-Behnken设计表（默认3水平）
    if flag == 2:
        print("设计方法： Box-Behnken设计")
        design = bbdesign(n=param, center=True) + 1
    # design3: 生成4个因子的完全正交表

    return design


def getCapacity(param, flag):
    # 初始化capacity
    capacity = {}
    for key in param.keys():
        if key not in capacity:
            capacity[key] = []

    if flag == 3:
        index = []
        for arr in param.values():
            index.append(len(arr))
    else:
        index = len(param)
    design = getDesign(index, flag)
    print(design)

    for arr in design:
        s = 0
        for key in capacity:
            capacity[key].append(param[key][int(arr[s])])
            s = s + 1

    return capacity


def add_height(params, substep, state_history, previous_state):
    sunlight = 0.5 if params['sunlight'] == 'low' else 1
    water = 0.5 if params['water'] == 'low' else 1
    temperature = 0.5 if params['temperature'] == 'low' else 1

    value = sunlight * water * temperature
    return {'height_add_value': value}


def add_leaf(params, substep, state_history, previous_state):
    sunlight = 0.5 if params['sunlight'] == 'low' else 1
    water = 0.5 if params['water'] == 'low' else 1
    temperature = 0.5 if params['temperature'] == 'low' else 1

    value = sunlight + water + temperature
    return {'leaf_add_value': value}


def update_height(params, substep, state_history, previous_state, policy_input):
    add_value = policy_input['height_add_value']
    height = previous_state['height']
    new_height = height + add_value
    # print(new_height)
    return 'height', new_height


def update_leaf_num(params, substep, state_history, previous_state, policy_input):
    add_value = policy_input['leaf_add_value']
    leaf_num = previous_state['leaf_num']
    new_leaf_num = leaf_num + add_value
    return 'leaf_num', new_leaf_num


psubs = [
    {
        'policies': {
            'height_add_value': add_height,
            'leaf_add_value': add_leaf
        },
        'variables': {
            'height': update_height,
            'leaf_num': update_leaf_num,
        }
    }
]


def run():
    del configs[:]
    capacity_param = {key: value for key, value in params.items() if
                      len(key) >= 5}
    capacity = getCapacity(capacity_param, 1)
    pprint(capacity)


if __name__ == "__main__":
    multiprocessing.freeze_support()
    run()
